[OMNIML-5613] Quantize ResNet residual adds in torch ONNX example - #2024

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Small, focused change (+90/-2) that adds per-block residual-add quantizers to the timm ResNet path of the torch→ONNX example, with CHANGELOG + README updates and an ONNX-level assertion in the existing example test. No licensing concerns; no prompt-injection content in the PR metadata. A few things worth addressing before merge:

1. Duplicated calibration logic / extra data pass. _add_resnet_residual_quantizers re-implements the enable_calib → forward-loop → load_calib_amax dance that _calibrate_uncalibrated_quantizers (same file) already performs, and adds a second full pass over the calibration set for ResNet (the first happens inside quantize_model for the FP8 Conv overrides). If the residual quantizer were attached as input_quantizer on the activation module and created before quantize_model, both the calibration and the dead-quantizer guard would come for free from the existing helpers. Even more idiomatic: modelopt already supports this via QuantModuleRegistry.register({nn.ReLU: "nn.ReLU"})(QuantInputBase) (exactly what modelopt/torch/quantization/nn/modules/quant_activations.py does for nn.LeakyReLU) plus a {"parent_class": "nn.ReLU", "quantizer_name": "*input_quantizer"} config entry, which gets calibrated by mtq.quantize's forward loop and is visible to mtq.print_quant_summary / modelopt state. Please either reuse one of these paths or note in the PR/comment why the manual hook is needed.

2. Residual quantizers bypass the file's own amax<=0/NaN guard, and load_calib_amax() is strict. _disable_dead_quantizers only inspects input_quantizer/output_quantizer/weight_quantizer, and it runs inside quantize_model — i.e. before these quantizers exist. A residual quantizer that calibrates to amax == 0 (or NaN) will therefore reach the FP8 exporter, which is precisely the scale = 448 / amax division the guard exists to prevent. Also, unlike _calibrate_uncalibrated_quantizers, this code calls load_calib_amax() without strict=False, so any block that didn't see data raises.

3. auto mode picks the residual format from the search space, not the search result. After mtq.auto_quantize, each block's actual format is known (e.g. block.conv3.input_quantizer._num_bits); deriving num_bits from the union of requested formats can give an FP8 residual Q/DQ next to an INT8-quantized block (or the reverse), which is what the rest of this file goes to some length to avoid for TRT.

Minor: the two new functions are the only helpers in this file without docstrings; and assert len(residual_adds) == 16 asserts that every Add in the exported graph is a residual add, which will break confusingly if the exporter ever emits an unrelated Add — consider filtering to the 16 residual adds (e.g. by producer/consumer pattern) before the count assertion.

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 223-239: Update _add_resnet_residual_quantizers and the
surrounding auto-quantization flow so residual quantizers are installed and
configured before mtq.auto_quantize() runs. For auto mode, include these
residual quantizer modules in every candidate format configuration used by the
search, ensuring their forced INT8/FP8 precision is scored and counted toward
the effective-bits constraint while preserving the existing non-auto behavior.

Actionable comments posted: 4

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 363-392: Update _finalize_resnet_quantizers so the int8/fp8 path
only accesses quantizer attributes when they exist, using safe getattr checks
for block.conv1, model.fc, model.global_pool, and related quantizers. Preserve
enabling/disabling behavior for present quantizers, and avoid AttributeError
when _prepare_resnet_quantizers skipped setup or a block variant lacks these
attributes.

In `@modelopt/onnx/export/fp8_exporter.py`:
- Line 34: Update the weight-channel threshold logic near the FP8
minimum-channel constant and the relevant convolution export check to account
for the convolution’s group count, matching the torch-side gate. Base the
decision on the effective total input channels rather than per-group
weight_input.values.shape[1], so grouped and depthwise convolutions consistently
enable or skip activation and weight Q/DQ paths.

In `@modelopt/onnx/utils.py`:
- Around line 1471-1481: The Cast-folding condition in the loop over
onnx_model.graph.node should not reject Cast-to-FLOAT nodes solely because
tensor_types lacks node.input[0]. Treat missing input type metadata as eligible
for folding, or resolve the producer output type before skipping; preserve the
existing FLOAT16 exclusion when the type is explicitly known.
- Around line 1524-1534: Replace the hard ValueError in the Q-consumer
validation with graceful skipping of that initializer, and emit a warning
identifying q_node.name and the observed scale dtype. Ensure
get_onnx_bytes_and_metadata continues exporting when a Q scale is a graph input
or remains non-FP16, while preserving FP16 conversion for valid Q scales.

---

Nitpick comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 232-259: Update _calibrate_new_quantizers to track quantizers
based on their pre-existing _if_quant state rather than is_enabled when building
enabled_quantizers. Restore quantization only for quantizers whose _if_quant
state was originally active, preserving intentionally disabled quantizers
through the finally block.

In `@modelopt/torch/_deploy/utils/torch_onnx.py`:
- Line 48: Promote _convert_q_data_initializers_to_fp16 in
modelopt/onnx/utils.py to a public helper name, add that name to the module’s
__all__, and update the import and call sites in the ONNX export flow around
_convert_q_data_initializers_to_fp16 accordingly. Preserve the existing
sequencing and validation behavior after the fold guard.

In `@tests/examples/torch_onnx/test_torch_quant_to_onnx.py`:
- Around line 35-62: Harden _assert_residual_adds_are_quantized against missing
graph producers by replacing direct producers[...] accesses with
producers.get(...) and explicit assertions that identify the missing input or
node. Apply the same pattern to the MatMul/ReduceMean lookup paths around the
next(...) calls, asserting the expected producer exists before dereferencing it
while preserving the current validation behavior.
- Around line 64-116: Add concise comments immediately before each hard-coded
topology-count assertion in the quantization test, explaining how the expected
values derive from the ResNet-50 Q/DQ structure and mode-specific behavior.
Cover the activation quantizer counts, DQ fanout counts, and int8 weight
quantizer count, including details such as the number of convolution weights and
disabled fully connected weight quantization; leave the assertions unchanged.

In `@tests/unit/onnx/test_fold_casts.py`:
- Around line 86-97: In `_initializer_q_model`, replace the redundant
`input_dtype` and `output_dtype` calculations with one shared dtype variable
derived from `opset`, and use it when constructing the primary `y` output
instead of creating it as `FLOAT16` and mutating it afterward. Remove the
subsequent `outputs[0].type.tensor_type.elem_type` assignment while preserving
the existing opset-dependent dtype.

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelopt/onnx/export/fp8_exporter.py`:
- Around line 196-197: Update the input-stem skip around the weight
dequantization pass so it applies only when the export is identified as a
ResNet, rather than using a three-channel input as the model-family
discriminator. Propagate and check an explicit ResNet/export marker alongside
node.inputs[0].name and weight_input.values.shape, while preserving the skip for
ResNet input stems and normal FP8 DQ restoration for non-ResNet models.

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelopt_recipes/README.md`:
- Around line 79-84: Update the “Choosing where to look” guidance in README.md
to add a timm-specific lookup step before the general fallback, directing
readers to timm/<architecture>/ for architecture-specific recipes. Preserve the
existing huggingface and general guidance while ensuring timm recipes are
included in the selection flow.

---

Nitpick comments:
In `@modelopt/torch/quantization/plugins/timm.py`:
- Around line 16-25: Add a module-level __all__ declaration in timm.py
containing only is_resnet_quantization_supported, preserving the intended public
API for wildcard imports.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Re-review of the reworked (recipe-driven) implementation. The approach changed substantially since the last round, and most of the earlier findings are genuinely resolved:

Resolved

• Duplicated calibration / second data pass (critical): the manual enable_calib → forward-loop → load_calib_amax block is gone. _add_resnet_residual_quantizers now runs as a CUSTOM_POST_CONVERSION_PLUGINS callback, i.e. after _replace_quant_module but before set_quantizer_by_cfg (see modelopt/torch/quantization/conversion.py::replace_quant_module), so the recipe's '*residual_quantizer' entry configures it and mtq.quantize's own forward loop calibrates it. The quantizers are also visible to print_quant_summary and modelopt state. Good fix.
• Strict load_calib_amax() (critical): moot — standard calibration path is used now.
• Auto-mode residual format inferred from the search space (critical): replaced by an explicit design — the residual quantizer is pinned to FP8 via fixed_quantization_config in the AutoQuantize recipe, documented in the README and PR body. Acceptable as an intentional choice (see one follow-up below).
• Test asserted "every Add is residual" (minor): the test now selects residual adds structurally by their single Relu consumer before the count assertion.

Still open / new

1. _disable_dead_quantizers still only inspects input_quantizer / output_quantizer / weight_quantizer, so a residual_quantizer that calibrates to amax == 0/NaN reaches the FP8 exporter's scale = 448 / amax — the exact case that guard was added for. The explicit NaN/non-positive check you described in the previous round no longer exists anywhere in the new code, so this half of the earlier comment regressed (one-line fix).
2. tests/examples/torch_onnx/test_torch_quant_to_onnx.py runs resnet50 + auto with --trt_build, but the PR body says the AutoQuantize FP8/INT8 mix only builds on Blackwell+. Please confirm how this test is expected to pass on the CI GPUs.
3. The recipe → mtq.auto_quantize kwargs translation duplicates examples/hf_ptq/hf_ptq.py::_mtq_inputs_from_auto_quantize_config, and the copy silently drops score_size (the new recipe sets score_size: 128 but the example uses --num_score_steps), cost_excluded_layers, and top-level candidate_formats.
4. --auto_quantization_formats is silently overridden for ResNet in auto mode with no warning and no opt-out (--recipe is rejected with --quantize_mode=auto).
5. Nits: the two new helpers are the only ones in this module without docstrings (the rationale docstring added in the previous round didn't survive the rewrite); "residual_quantizer" in block.downsample._modules could use hasattr; modelopt_recipes/README.md still doesn't mention the new top-level timm/ tree (earlier bot comment unaddressed); timm/resnet/ptq/{fp8,int8}.yaml are verbatim copies of configs/ptq/presets/model/{fp8,int8}.yaml plus one entry, so they will drift if the presets change.

No licensing concerns (new YAMLs carry the standard NVIDIA SPDX header); no prompt-injection content in the PR metadata.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

All previously-flagged critical items are resolved (residual quantizers now go through the standard calibration flow via CUSTOM_POST_CONVERSION_PLUGINS, _disable_dead_quantizers covers residual_quantizer, residual adds are selected structurally in the test, --recipe is authoritative instead of silently overriding flags, and the AutoQuantize mapping no longer drops score_size/cost_excluded_layers/candidate_formats). Design-wise this reuses the existing recipe/preset/plugin machinery rather than inventing a new one — good.

Three things I'd like addressed before merge:

1. _match_qformat can never match a recipe-derived config, so every recipe path silently skips the TRT post-processing. QUANT_CFG_CHOICES entries are model_dump(exclude_unset=True) (sparse: {"quantizer_name": "*", "enable": false}), while _match_qformat compares them against recipe.quantize.model_dump() (dense: parent_class: None, cfg: None, plus ~20 QuantizerAttributeConfig defaults). entry in config["quant_cfg"] is therefore always False, so _recipe_qformats() returns an empty set and _prepare_auto_quantize_format() never adds the Conv2d override. Concretely: --recipe=timm/resnet/ptq/fp8 never runs _disable_low_channel_conv_input_quantizers, i.e. the raw-RGB conv1 keeps its FP8 input quantizer — the exact Blackwell "no implementation for TRT_FP8QuantizeLinear" failure the helper exists to prevent. CI on Ada hides this. Same class of bug for recipe-supplied mxfp8/nvfp4 AutoQuantize candidates (no Conv2d→FP8 override, and _disable_high_rank_input_quantizers skipped).

2. ResNet + auto is silently dropped from a previously-supported/tested configuration (README ✅ → blank, test pytest.skip("AutoQuantize is not supported for ResNet")). The flag-based auto path for resnet50 with --trt_build was passing before this PR; if the reason is the now-unconditional residual-quantizer registration, register the plugin only when residual quantization is actually requested, and otherwise please explain the regression in the PR body/CHANGELOG.

3. New recipe YAMLs carry a truncated license header (2-line SPDX only) whereas every other modelopt_recipes/**.yaml uses the full LICENSE_HEADER Apache text. Not auto-enforced (the insert-license hook only covers python/shell/c), so it needs a manual fix.

Also: there is no unit coverage for _match_qformat / _recipe_qformats / _mtq_inputs_from_auto_quantize_config — a small pure-Python test asserting that the new timm/resnet/ptq/fp8.yaml resolves to {"fp8"} would have caught (1) immediately.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Re-review (8 files, +448/-144). Design-review protocol applied first: this PR does not introduce a new subsystem — the residual quantizer is installed through the existing CUSTOM_POST_CONVERSION_PLUGINS hook (invoked by replace_quant_module before set_quantizer_by_cfg, so the recipe configures it and mtq.quantize's own forward loop calibrates it), and the new YAMLs slot into the existing modelopt_recipes/ catalog with $import units. That's reuse of the in-repo pattern, not a second one. Licensing is fine: the two new recipe YAMLs now carry the verbatim LICENSE_HEADER text. No prompt-injection content in the PR metadata.

Previous blockers — resolved

• _match_qformat could never match a recipe-derived config, so --recipe silently skipped TRT post-processing: the reverse-preset-matching helpers are gone. _disable_low_channel_fp8_conv_input_quantizers and _disable_high_rank_input_quantizers now run unconditionally and decide from live quantizer state (num_bits == (4, 3), block_sizes), so --recipe=timm/resnet/ptq/fp8 does disable the raw-RGB conv1 input quantizer. Confirmed the loader converts eXmY → tuples, so the (4,3)/(2,1) comparisons are sound.
• Truncated SPDX-only header on the new YAMLs: fixed, matches LICENSE_HEADER.
• modelopt_recipes/README.md didn't mention the timm/ tree: fixed (catalog + selection guidance + placement), plus a timm/resnet/ptq/README.md documenting the delta.
• AutoQuantize mapping dropped score_size/cost_excluded_layers/candidate_formats: all three are mapped now, and test_auto_quantize_recipe_mapping + a new e2e recipe test cover the top-level-candidates shape (previously the module_search_spaces=[] failure case).
• Residual plugin registered unconditionally: now gated on the recipe actually enabling *residual_quantizer, and _disable_dead_quantizers covers residual_quantizer.

Also verified the removal of _calibrate_uncalibrated_quantizers is safe: forcing algorithm = "max" for mxfp8/nvfp4 calibrates the FP8 Conv overrides in the normal pass, and finish_stats_collection skips amax loading for dynamic/MX quantizers.

Why a nudge rather than an approve — see the reason field: the remaining items are product/behavior judgment calls plus a few small robustness gaps, not a re-flag of the fixed bugs.

LGTM, let's move efforts to make this a more generic workflow to a follow-up PR. Thanks!